"""
数据存储服务
支持SQLite、InfluxDB和MySQL多种存储后端
"""

import sqlite3
import logging
import threading
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from contextlib import contextmanager

try:
    from influxdb_client import InfluxDBClient, Point, WritePrecision
    from influxdb_client.client.write_api import SYNCHRONOUS
    INFLUXDB_AVAILABLE = True
except ImportError:
    INFLUXDB_AVAILABLE = False

try:
    import pymysql
    MYSQL_AVAILABLE = True
except ImportError:
    MYSQL_AVAILABLE = False

from core.base_driver import TagValue

class StorageService:
    """数据存储服务"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.logger = logging.getLogger(self.__class__.__name__)
        
        # 存储后端
        self.sqlite_enabled = config.get('sqlite', {}).get('enabled', True)
        self.influxdb_enabled = config.get('influxdb', {}).get('enabled', False)
        self.mysql_enabled = config.get('mysql', {}).get('enabled', False)
        
        # 初始化存储后端
        self.sqlite_conn = None
        self.influxdb_client = None
        self.mysql_pool = None
        
        # 线程锁
        self.lock = threading.RLock()
        
        self._init_storage_backends()
    
    def _init_storage_backends(self):
        """初始化存储后端"""
        if self.sqlite_enabled:
            self._init_sqlite()
        
        if self.influxdb_enabled:
            self._init_influxdb()
        
        if self.mysql_enabled:
            self._init_mysql()
    
    def _init_sqlite(self):
        """初始化SQLite数据库"""
        try:
            sqlite_config = self.config.get('sqlite', {})
            db_path = sqlite_config.get('path', 'data/iot_collector.db')
            
            # 创建数据库目录
            import os
            os.makedirs(os.path.dirname(db_path), exist_ok=True)
            
            self.sqlite_conn = sqlite3.connect(
                db_path, 
                check_same_thread=False,
                timeout=30
            )
            self.sqlite_conn.row_factory = sqlite3.Row
            
            # 创建表结构
            self._create_sqlite_tables()
            
            self.logger.info("SQLite数据库初始化成功")
            
        except Exception as e:
            self.logger.error(f"SQLite数据库初始化失败: {e}")
            self.sqlite_enabled = False
    
    def _create_sqlite_tables(self):
        """创建SQLite表结构"""
        with self.sqlite_conn:
            # 实时数据表
            self.sqlite_conn.execute('''
                CREATE TABLE IF NOT EXISTS realtime_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    device_name TEXT NOT NULL,
                    tag_name TEXT NOT NULL,
                    value TEXT,
                    timestamp DATETIME NOT NULL,
                    quality TEXT DEFAULT 'GOOD',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建索引
            self.sqlite_conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_realtime_device_tag 
                ON realtime_data (device_name, tag_name)
            ''')
            self.sqlite_conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_realtime_timestamp 
                ON realtime_data (timestamp)
            ''')
            
            # 历史数据表
            self.sqlite_conn.execute('''
                CREATE TABLE IF NOT EXISTS historical_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    device_name TEXT NOT NULL,
                    tag_name TEXT NOT NULL,
                    value TEXT,
                    timestamp DATETIME NOT NULL,
                    quality TEXT DEFAULT 'GOOD',
                    archived_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建索引
            self.sqlite_conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_historical_device_tag_time 
                ON historical_data (device_name, tag_name, timestamp)
            ''')
            
            # 设备状态表
            self.sqlite_conn.execute('''
                CREATE TABLE IF NOT EXISTS device_status (
                    device_name TEXT PRIMARY KEY,
                    connection_status TEXT,
                    last_communication DATETIME,
                    error_count INTEGER DEFAULT 0,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 报警记录表
            self.sqlite_conn.execute('''
                CREATE TABLE IF NOT EXISTS alarm_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    device_name TEXT NOT NULL,
                    tag_name TEXT NOT NULL,
                    alarm_type TEXT NOT NULL,
                    alarm_level TEXT NOT NULL,
                    message TEXT,
                    value TEXT,
                    timestamp DATETIME NOT NULL,
                    acknowledged BOOLEAN DEFAULT FALSE,
                    acknowledged_by TEXT,
                    acknowledged_at DATETIME,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建索引
            self.sqlite_conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_alarm_device_time 
                ON alarm_records (device_name, timestamp)
            ''')
            self.sqlite_conn.execute('''
                CREATE INDEX IF NOT EXISTS idx_alarm_level 
                ON alarm_records (alarm_level)
            ''')
    
    def _init_influxdb(self):
        """初始化InfluxDB连接"""
        if not INFLUXDB_AVAILABLE:
            self.logger.error("influxdb-client库未安装")
            self.influxdb_enabled = False
            return
        
        try:
            influxdb_config = self.config.get('influxdb', {})
            
            self.influxdb_client = InfluxDBClient(
                url=influxdb_config.get('url', 'http://localhost:8086'),
                token=influxdb_config.get('token', ''),
                org=influxdb_config.get('org', 'iot')
            )
            
            # 测试连接
            self.influxdb_client.ping()
            
            self.influx_bucket = influxdb_config.get('bucket', 'iot_data')
            self.influx_org = influxdb_config.get('org', 'iot')
            
            self.write_api = self.influxdb_client.write_api(write_options=SYNCHRONOUS)
            self.query_api = self.influxdb_client.query_api()
            
            self.logger.info("InfluxDB连接初始化成功")
            
        except Exception as e:
            self.logger.error(f"InfluxDB连接初始化失败: {e}")
            self.influxdb_enabled = False
    
    def _init_mysql(self):
        """初始化MySQL连接"""
        if not MYSQL_AVAILABLE:
            self.logger.error("pymysql库未安装")
            self.mysql_enabled = False
            return
        
        try:
            mysql_config = self.config.get('mysql', {})
            
            self.mysql_config = {
                'host': mysql_config.get('host', 'localhost'),
                'port': mysql_config.get('port', 3306),
                'user': mysql_config.get('username', 'root'),
                'password': mysql_config.get('password', ''),
                'database': mysql_config.get('database', 'iot_collector'),
                'charset': 'utf8mb4',
                'autocommit': True
            }
            
            # 测试连接
            with self._get_mysql_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("SELECT 1")
            
            self._create_mysql_tables()
            
            self.logger.info("MySQL连接初始化成功")
            
        except Exception as e:
            self.logger.error(f"MySQL连接初始化失败: {e}")
            self.mysql_enabled = False
    
    @contextmanager
    def _get_mysql_connection(self):
        """获取MySQL连接"""
        conn = pymysql.connect(**self.mysql_config)
        try:
            yield conn
        finally:
            conn.close()
    
    def _create_mysql_tables(self):
        """创建MySQL表结构"""
        with self._get_mysql_connection() as conn:
            with conn.cursor() as cursor:
                # 实时数据表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS realtime_data (
                        id BIGINT AUTO_INCREMENT PRIMARY KEY,
                        device_name VARCHAR(100) NOT NULL,
                        tag_name VARCHAR(100) NOT NULL,
                        value TEXT,
                        timestamp DATETIME(3) NOT NULL,
                        quality VARCHAR(20) DEFAULT 'GOOD',
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        INDEX idx_device_tag (device_name, tag_name),
                        INDEX idx_timestamp (timestamp)
                    ) ENGINE=InnoDB
                ''')
                
                # 历史数据表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS historical_data (
                        id BIGINT AUTO_INCREMENT PRIMARY KEY,
                        device_name VARCHAR(100) NOT NULL,
                        tag_name VARCHAR(100) NOT NULL,
                        value TEXT,
                        timestamp DATETIME(3) NOT NULL,
                        quality VARCHAR(20) DEFAULT 'GOOD',
                        archived_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        INDEX idx_device_tag_time (device_name, tag_name, timestamp)
                    ) ENGINE=InnoDB
                ''')
    
    def save_tag_values(self, tag_values: List[TagValue]):
        """保存标签值数据"""
        if not tag_values:
            return
        
        with self.lock:
            # 保存到SQLite
            if self.sqlite_enabled:
                self._save_to_sqlite(tag_values)
            
            # 保存到InfluxDB
            if self.influxdb_enabled:
                self._save_to_influxdb(tag_values)
            
            # 保存到MySQL
            if self.mysql_enabled:
                self._save_to_mysql(tag_values)
    
    def _save_to_sqlite(self, tag_values: List[TagValue]):
        """保存数据到SQLite"""
        try:
            with self.sqlite_conn:
                cursor = self.sqlite_conn.cursor()
                
                # 清理旧的实时数据（保留最近1小时）
                cleanup_time = datetime.now() - timedelta(hours=1)
                cursor.execute(
                    "DELETE FROM realtime_data WHERE timestamp < ?",
                    (cleanup_time,)
                )
                
                # 插入新数据
                for tag_value in tag_values:
                    cursor.execute('''
                        INSERT OR REPLACE INTO realtime_data 
                        (device_name, tag_name, value, timestamp, quality)
                        VALUES (?, ?, ?, ?, ?)
                    ''', (
                        tag_value.device_name,
                        tag_value.tag_name,
                        str(tag_value.value),
                        tag_value.timestamp,
                        tag_value.quality
                    ))
                
                self.sqlite_conn.commit()
                
        except Exception as e:
            self.logger.error(f"保存数据到SQLite失败: {e}")
    
    def _save_to_influxdb(self, tag_values: List[TagValue]):
        """保存数据到InfluxDB"""
        try:
            points = []
            for tag_value in tag_values:
                point = Point(tag_value.tag_name) \
                    .tag("device", tag_value.device_name) \
                    .tag("quality", tag_value.quality) \
                    .field("value", tag_value.value) \
                    .time(tag_value.timestamp, WritePrecision.NS)
                points.append(point)
            
            self.write_api.write(
                bucket=self.influx_bucket,
                org=self.influx_org,
                record=points
            )
            
        except Exception as e:
            self.logger.error(f"保存数据到InfluxDB失败: {e}")
    
    def _save_to_mysql(self, tag_values: List[TagValue]):
        """保存数据到MySQL"""
        try:
            with self._get_mysql_connection() as conn:
                with conn.cursor() as cursor:
                    # 批量插入数据
                    data = [
                        (
                            tag_value.device_name,
                            tag_value.tag_name,
                            str(tag_value.value),
                            tag_value.timestamp,
                            tag_value.quality
                        )
                        for tag_value in tag_values
                    ]
                    
                    cursor.executemany('''
                        INSERT INTO realtime_data 
                        (device_name, tag_name, value, timestamp, quality)
                        VALUES (%s, %s, %s, %s, %s)
                        ON DUPLICATE KEY UPDATE
                        value = VALUES(value),
                        timestamp = VALUES(timestamp),
                        quality = VALUES(quality)
                    ''', data)
                    
        except Exception as e:
            self.logger.error(f"保存数据到MySQL失败: {e}")
    
    def get_latest_values(self, device_name: str = None, tag_name: str = None) -> List[Dict[str, Any]]:
        """获取最新数据"""
        if self.sqlite_enabled:
            return self._get_latest_from_sqlite(device_name, tag_name)
        elif self.mysql_enabled:
            return self._get_latest_from_mysql(device_name, tag_name)
        else:
            return []
    
    def _get_latest_from_sqlite(self, device_name: str = None, tag_name: str = None) -> List[Dict[str, Any]]:
        """从SQLite获取最新数据"""
        try:
            cursor = self.sqlite_conn.cursor()
            
            sql = "SELECT * FROM realtime_data WHERE 1=1"
            params = []
            
            if device_name:
                sql += " AND device_name = ?"
                params.append(device_name)
            
            if tag_name:
                sql += " AND tag_name = ?"
                params.append(tag_name)
            
            sql += " ORDER BY timestamp DESC LIMIT 1000"
            
            cursor.execute(sql, params)
            rows = cursor.fetchall()
            
            return [dict(row) for row in rows]
            
        except Exception as e:
            self.logger.error(f"从SQLite获取数据失败: {e}")
            return []
    
    def _get_latest_from_mysql(self, device_name: str = None, tag_name: str = None) -> List[Dict[str, Any]]:
        """从MySQL获取最新数据"""
        try:
            with self._get_mysql_connection() as conn:
                with conn.cursor(pymysql.cursors.DictCursor) as cursor:
                    sql = "SELECT * FROM realtime_data WHERE 1=1"
                    params = []
                    
                    if device_name:
                        sql += " AND device_name = %s"
                        params.append(device_name)
                    
                    if tag_name:
                        sql += " AND tag_name = %s"
                        params.append(tag_name)
                    
                    sql += " ORDER BY timestamp DESC LIMIT 1000"
                    
                    cursor.execute(sql, params)
                    return cursor.fetchall()
                    
        except Exception as e:
            self.logger.error(f"从MySQL获取数据失败: {e}")
            return []
    
    def get_historical_data(self, device_name: str, tag_name: str, 
                          start_time: datetime, end_time: datetime) -> List[Dict[str, Any]]:
        """获取历史数据"""
        if self.influxdb_enabled:
            return self._get_historical_from_influxdb(device_name, tag_name, start_time, end_time)
        elif self.mysql_enabled:
            return self._get_historical_from_mysql(device_name, tag_name, start_time, end_time)
        else:
            return []
    
    def query_historical_data(self, device_name: str, tag_name: str = None, 
                            start_time: datetime = None, end_time: datetime = None) -> List[Dict[str, Any]]:
        """查询历史数据（兼容方法）"""
        try:
            if not start_time:
                start_time = datetime.now() - timedelta(hours=24)
            if not end_time:
                end_time = datetime.now()
            
            if tag_name:
                return self.get_historical_data(device_name, tag_name, start_time, end_time)
            else:
                # 如果没有指定标签，返回设备的所有历史数据
                if self.sqlite_enabled and self.sqlite_conn:
                    with self.lock:
                        cursor = self.sqlite_conn.cursor()
                        cursor.execute('''
                            SELECT device_name, tag_name, value, timestamp, quality
                            FROM realtime_data 
                            WHERE device_name = ? AND timestamp BETWEEN ? AND ?
                            ORDER BY timestamp DESC
                            LIMIT 1000
                        ''', (device_name, start_time, end_time))
                        
                        rows = cursor.fetchall()
                        cursor.close()
                        
                        return [
                            {
                                'device_name': row[0],
                                'tag_name': row[1],
                                'value': row[2],
                                'timestamp': row[3],
                                'quality': row[4]
                            }
                            for row in rows
                        ]
                return []
        except Exception as e:
            self.logger.error(f"查询历史数据失败: {e}")
            return []
    
    def _get_historical_from_influxdb(self, device_name: str, tag_name: str,
                                    start_time: datetime, end_time: datetime) -> List[Dict[str, Any]]:
        """从InfluxDB获取历史数据"""
        try:
            query = f'''
                from(bucket: "{self.influx_bucket}")
                |> range(start: {start_time.isoformat()}Z, stop: {end_time.isoformat()}Z)
                |> filter(fn: (r) => r["_measurement"] == "{tag_name}")
                |> filter(fn: (r) => r["device"] == "{device_name}")
                |> filter(fn: (r) => r["_field"] == "value")
            '''
            
            tables = self.query_api.query(query, org=self.influx_org)
            
            result = []
            for table in tables:
                for record in table.records:
                    result.append({
                        'device_name': record.values.get('device'),
                        'tag_name': record.get_measurement(),
                        'value': record.get_value(),
                        'timestamp': record.get_time(),
                        'quality': record.values.get('quality', 'GOOD')
                    })
            
            return result
            
        except Exception as e:
            self.logger.error(f"从InfluxDB获取历史数据失败: {e}")
            return []
    
    def _get_historical_from_mysql(self, device_name: str, tag_name: str,
                                 start_time: datetime, end_time: datetime) -> List[Dict[str, Any]]:
        """从MySQL获取历史数据"""
        try:
            with self._get_mysql_connection() as conn:
                with conn.cursor(pymysql.cursors.DictCursor) as cursor:
                    cursor.execute('''
                        SELECT * FROM historical_data 
                        WHERE device_name = %s AND tag_name = %s 
                        AND timestamp BETWEEN %s AND %s
                        ORDER BY timestamp ASC
                    ''', (device_name, tag_name, start_time, end_time))
                    
                    return cursor.fetchall()
                    
        except Exception as e:
            self.logger.error(f"从MySQL获取历史数据失败: {e}")
            return []
    
    def save_device_status(self, device_name: str, status: Dict[str, Any]):
        """保存设备状态"""
        if self.sqlite_enabled:
            self._save_device_status_sqlite(device_name, status)
    
    def _save_device_status_sqlite(self, device_name: str, status: Dict[str, Any]):
        """保存设备状态到SQLite"""
        try:
            with self.sqlite_conn:
                self.sqlite_conn.execute('''
                    INSERT OR REPLACE INTO device_status 
                    (device_name, connection_status, last_communication, error_count, updated_at)
                    VALUES (?, ?, ?, ?, ?)
                ''', (
                    device_name,
                    status.get('connection_status'),
                    status.get('last_communication'),
                    status.get('error_count', 0),
                    datetime.now()
                ))
                
        except Exception as e:
            self.logger.error(f"保存设备状态失败: {e}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取存储统计信息"""
        stats = {
            'backends': {
                'sqlite': self.sqlite_enabled,
                'influxdb': self.influxdb_enabled,
                'mysql': self.mysql_enabled
            },
            'data_points': 0,
            'devices': 0,
            'tags': 0
        }
        
        try:
            if self.sqlite_enabled and self.sqlite_conn:
                with self.lock:
                    cursor = self.sqlite_conn.cursor()
                    
                    # 获取数据点总数
                    cursor.execute("SELECT COUNT(*) FROM realtime_data")
                    stats['data_points'] = cursor.fetchone()[0]
                    
                    # 获取设备数量
                    cursor.execute("SELECT COUNT(DISTINCT device_name) FROM realtime_data")
                    stats['devices'] = cursor.fetchone()[0]
                    
                    # 获取标签数量
                    cursor.execute("SELECT COUNT(DISTINCT device_name || '.' || tag_name) FROM realtime_data")
                    stats['tags'] = cursor.fetchone()[0]
                    
                    cursor.close()
        except Exception as e:
            self.logger.error(f"获取存储统计信息失败: {e}")
        
        return stats

    def close(self):
        """关闭存储连接"""
        if self.sqlite_conn:
            self.sqlite_conn.close()
        
        if self.influxdb_client:
            self.influxdb_client.close()
        
        self.logger.info("存储服务已关闭") 